import numpy as np
import matplotlib.pyplot as plt
from matplotlib import rcParams


rcParams['font.family'] = 'sans-serif'
rcParams['font.sans-serif'] = ['Arial']


# Read data from file skipping lines starting with '@'
data = []
with open('wt_pro8_trajout/fluctpro.agr', 'r') as file:
    for line in file:
        if not line.startswith('@'):
            data.append([float(x) for x in line.split()])

# Convert data to numpy array
data_array = np.array(data)

# Extract residue numbers and RMSF values
residue_numbers = data_array[:, 0].astype(int)
rmsf_values = data_array[:, 1]

# Normalize RMSF values to [0, 1] for colormap
normalized_rmsf = (rmsf_values - rmsf_values.min()) / (rmsf_values.max() - rmsf_values.min())

# Get colors from the viridis colormap
cmap = plt.get_cmap('BuPu')
colors = cmap(normalized_rmsf)

# Create a PyMOL script
with open('WT_color_rmsf_2.pml', 'w') as pymol_script:
    for res_num, color in zip(residue_numbers, colors):
        color_str = f'{color[0]:.2f}, {color[1]:.2f}, {color[2]:.2f}'
        pymol_script.write(f'set_color res_color_{res_num}, [{color_str}]\n')
        pymol_script.write(f'color res_color_{res_num}, resi {res_num}\n')
    pymol_script.write('show cartoon\n')  # or any representation you prefer
    pymol_script.write('bg_color white\n')

print("PyMOL script 'WT_color_rmsf.pml' has been created. Load this script in PyMOL to visualize the RMSF values.")

# Create and save the color scale (color bar) as a PNG file
fig, ax = plt.subplots(figsize=(8, 1))
fig.subplots_adjust(bottom=0.5)

# Create a color bar with the viridis colormap, using the RMSF values as the range
norm = plt.Normalize(vmin=rmsf_values.min(), vmax=rmsf_values.max())
cbar = plt.colorbar(plt.cm.ScalarMappable(norm=norm, cmap=cmap), cax=ax, orientation='horizontal')
cbar.set_label('RMSF Fluctuation, Å', fontsize=18)

# Save the color bar to a PNG file
plt.savefig('RMSF_color_scale.png')

print("Color scale has been saved to 'color_scale.png'.")
